import jsonlines
from battery_qa_generator.qa_generator import QAGenerator

class QAmain:
    def __init__(self, cycle_data, step_data, record_data, cycle_group_indices, charge_discharge_rates, charge_discharge_currents, reactant_mass):
        self.cycle_data = cycle_data
        self.step_data = step_data
        self.record_data = record_data
        self.cycle_group_indices = cycle_group_indices
        self.charge_discharge_rates = charge_discharge_rates
        self.charge_discharge_currents = charge_discharge_currents
        self.reactant_mass = reactant_mass

    def main(self):
        generator = QAGenerator(
            self.cycle_data, 
            self.step_data, 
            self.record_data,
            self.cycle_group_indices, 
            self.charge_discharge_rates, 
            self.charge_discharge_currents,
            self.reactant_mass,
            )
        # 生成问答对
        qa_pairs = generator.generate_all()
        #print(qa_pairs)
        # 保存结果
        with open("batllm_data_oxide_before100_lqy1.jsonl", 'a', encoding="utf-8") as file:
            writer = jsonlines.Writer(file)
            for qa in qa_pairs:
                writer.write({
                'input': qa['input'],
                'output': qa['output']
            })
        writer.close()